package com.shawn.integeration.ratelimiter.service.impl;

import com.shawn.integeration.log.util.IpUtils;
import com.shawn.integeration.ratelimiter.annonation.RateLimiter;
import com.shawn.integeration.ratelimiter.enums.LimitType;
import com.shawn.integeration.ratelimiter.service.IRateLimiterService;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.lang.reflect.Method;
import java.util.Collections;
import java.util.List;
import java.util.Objects;

/**
 * 默认限流实现
 *
 * @author: shawn
 * @Date: 2023/12/28 10:44
 * @Version 1.0.0
 */
@Slf4j
@Service
public class DefaultRateLimiterServiceImpl implements IRateLimiterService {

    @Resource
    private RedisTemplate<Object, Object> redisTemplate;
    @Resource
    private RedisScript<Long> limitScript;

    @Override
    public boolean canAccess(JoinPoint joinPoint, RateLimiter rateLimiter) {
        // 限流时间和限流次数
        int time = rateLimiter.time();
        int count = rateLimiter.count();
        // 获取限流 Key
        String combineKey = buildLimitKey(rateLimiter, joinPoint);

        // 执行 redis lua 脚本进行限流
        List<Object> keys = Collections.singletonList(combineKey);
        Long number = redisTemplate.execute(limitScript, keys, count, time);
        if (Objects.nonNull(number) && number.intValue() > count) {
            return false;
        }

        log.info("限制请求'{}',当前请求'{}',缓存key'{}'", count, number.intValue(), combineKey);
        return true;
    }

    /**
     * 构建限流 key, key = 限流前缀 + [请求IP] + 类名 + 方法名
     */
    public String buildLimitKey(RateLimiter rateLimiter, JoinPoint point) {
        StringBuilder stringBuffer = new StringBuilder(rateLimiter.key());
        if (rateLimiter.limitType() == LimitType.IP) {
            stringBuffer.append(IpUtils.getIpAddr()).append("-");
        }
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        Class<?> targetClass = method.getDeclaringClass();
        stringBuffer.append(targetClass.getName()).append("-").append(method.getName());
        return stringBuffer.toString();
    }
}
